import torch.nn as nn
import torch

class SPARSE(nn.Module):
    def __init__(self,params,N=2,std_w=1e-3,mode='double'):
        # Initial the parameter (SPARSE)
        super().__init__()
        data_shape = params
        self.mode = mode
        self.u = nn.Parameter((torch.randn(data_shape)-0.5)*2*std_w)
        self.v = nn.Parameter((torch.randn(data_shape)-0.5)*2*std_w)
        self.N = N



    def forward(self,*args):
        if self.mode == 'double':
            return self.u**self.N-self.v**self.N
        else:
            return self.u**self.N








